"""
智能体基类 - 基于LangGraph ReAct模式

使用LangGraph的create_react_agent构建智能体，支持：
1. 工具调用能力
2. 推理过程展示
3. 流式输出
4. 状态管理
"""

import asyncio
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, List, AsyncIterator, cast
from dataclasses import dataclass, field
from datetime import datetime

from langchain.chat_models import init_chat_model
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage
from langchain_core.tools import Tool
from langchain_core.runnables import RunnableConfig
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from pydantic import SecretStr

from logging_config import get_logger


@dataclass
class LLMConfig:
    """LLM配置类"""
    provider: str = "openai"
    model: str = "gpt-4"
    api_key: str = ""
    api_base: str = ""
    temperature: float = 0.1
    timeout: int = 60
    streaming: bool = True

    def __post_init__(self):
        """初始化后处理"""
        # 如果api_base为空，设置默认值
        if not self.api_base:
            if self.provider == "openai":
                self.api_base = "https://api.openai.com/v1"
            elif self.provider == "anthropic":
                self.api_base = "https://api.anthropic.com"


@dataclass
class AgentState:
    """智能体状态"""
    messages: List[BaseMessage]
    reasoning_steps: List[str]
    tool_calls: List[Dict[str, Any]]
    current_step: int
    final_answer: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


class BaseAgent(ABC):
    """智能体基类 - 基于LangGraph ReAct模式"""

    def __init__(self, name: str, description: str = "", llm_config: Optional[LLMConfig] = None):
        """
        初始化智能体

        Args:
            name: 智能体名称
            description: 智能体描述
            llm_config: LLM配置
        """
        self.name = name
        self.description = description

        # 初始化LLM配置，如果没有提供则从settings加载
        if llm_config is None:
            from config.settings import settings
            self.llm_config = LLMConfig(
                provider=settings.LLM_PROVIDER,
                model=settings.OPENAI_MODEL if settings.LLM_PROVIDER == "openai" else settings.ANTHROPIC_MODEL,
                api_key=settings.OPENAI_API_KEY if settings.LLM_PROVIDER == "openai" else settings.ANTHROPIC_API_KEY,
                api_base=settings.OPENAI_API_BASE if settings.LLM_PROVIDER == "openai" else "",
                temperature=0.1,
                timeout=60,
                streaming=True
            )
        else:
            self.llm_config = llm_config

        # 设置日志 - 使用统一的日志系统
        self.logger = get_logger(f"agent.{name}")

        # LLM实例（延迟初始化）
        self._llm = None

        # ReAct智能体（延迟初始化）
        self._react_agent = None

        # 检查点保存器（用于状态持久化）
        self._checkpointer = MemorySaver()

        # 工具列表
        self._tools = []
    
    @property
    def llm(self):
        """获取LLM实例（延迟初始化）"""
        if self._llm is None:
            self.logger.info(f"初始化LLM: provider={self.llm_config.provider}, model={self.llm_config.model}")
            self.logger.info(f"API配置: base_url={self.llm_config.api_base}, api_key={'***' if self.llm_config.api_key else 'None'}")

            try:
                self._llm = init_chat_model(
                    model_provider=self.llm_config.provider,
                    model=self.llm_config.model,
                    temperature=self.llm_config.temperature,
                    base_url=self.llm_config.api_base,
                    api_key=SecretStr(self.llm_config.api_key),
                    timeout=self.llm_config.timeout,
                    streaming=self.llm_config.streaming,
                )
                self.logger.info("LLM初始化成功")
            except Exception as e:
                self.logger.error(f"LLM初始化失败: {e}")
                raise
        return self._llm
    
    @property
    def react_agent(self):
        """获取ReAct智能体（延迟初始化）"""
        if self._react_agent is None:
            # 创建ReAct智能体
            self._react_agent = create_react_agent(
                model=self.llm,
                tools=self._tools,
                checkpointer=self._checkpointer,
            )
        return self._react_agent
    
    @abstractmethod
    def get_tools(self) -> List[Tool]:
        """
        获取智能体可用的工具列表（子类必须实现）

        Returns:
            工具列表
        """
        pass

    def _create_sync_wrapper(self, async_func):
        """为异步工具函数创建同步包装器"""
        import asyncio

        def sync_wrapper(*args, **kwargs):
            """同步包装器 - 支持嵌套事件循环"""
            tool_name = async_func.__name__
            self.logger.info(f"🔧 工具包装器被调用: {tool_name}, args={args}, kwargs={kwargs}")
            
            try:
                # 尝试获取当前运行的事件循环
                try:
                    loop = asyncio.get_running_loop()
                    self.logger.info(f"检测到运行中的事件循环，使用nest_asyncio")
                    # 如果有运行中的循环，使用nest_asyncio来允许嵌套
                    try:
                        import nest_asyncio
                        nest_asyncio.apply()
                        self.logger.info(f"nest_asyncio已应用，开始执行异步函数: {tool_name}")
                        result = asyncio.run(async_func(*args, **kwargs))
                        self.logger.info(f"✅ 异步函数执行成功: {tool_name}, 返回类型: {type(result)}, 长度: {len(str(result))}")
                        return result
                    except ImportError:
                        # 如果没有nest_asyncio，尝试使用run_until_complete
                        self.logger.warning("nest_asyncio未安装，尝试使用备用方法")
                        # 创建新的事件循环来执行
                        new_loop = asyncio.new_event_loop()
                        try:
                            result = new_loop.run_until_complete(async_func(*args, **kwargs))
                            self.logger.info(f"✅ 使用新事件循环执行成功: {tool_name}")
                            return result
                        finally:
                            new_loop.close()
                except RuntimeError as e:
                    # 没有运行中的事件循环，直接使用asyncio.run
                    self.logger.info(f"未检测到运行中的事件循环，直接使用asyncio.run")
                    result = asyncio.run(async_func(*args, **kwargs))
                    self.logger.info(f"✅ 直接执行成功: {tool_name}")
                    return result
            except Exception as e:
                # 记录详细错误信息
                self.logger.error(f"❌ 异步工具函数调用失败: {tool_name}", exc_info=True)
                self.logger.error(f"错误类型: {type(e).__name__}")
                self.logger.error(f"错误信息: {e}")
                self.logger.error(f"参数: args={args}, kwargs={kwargs}")
                return f"工具{tool_name}调用失败: {type(e).__name__}: {e}"

        return sync_wrapper
    
    @abstractmethod
    def get_system_prompt(self) -> str:
        """
        获取系统提示词（子类必须实现）
        
        Returns:
            系统提示词
        """
        pass
    
    def initialize_tools(self):
        """初始化工具"""
        if not self._tools:
            raw_tools = self.get_tools()
            # 为异步工具函数创建同步包装器
            wrapped_tools = []
            for tool in raw_tools:
                if asyncio.iscoroutinefunction(tool.func):
                    self.logger.info(f"检测到异步工具函数: {tool.name}，创建同步包装器")
                    sync_wrapper = self._create_sync_wrapper(tool.func)
                    # 创建新的Tool对象，使用同步包装器
                    wrapped_tool = Tool(
                        name=tool.name,
                        description=tool.description,
                        func=sync_wrapper
                    )
                    wrapped_tools.append(wrapped_tool)
                else:
                    wrapped_tools.append(tool)

            self._tools = wrapped_tools
            self.logger.info(f"初始化{len(self._tools)}个工具")
    
    async def stream_analyze(
        self,
        input_data: Any,
        thread_id: Optional[str] = None,
        **kwargs
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        流式分析方法（核心方法）

        Args:
            input_data: 输入数据
            thread_id: 线程ID（用于状态持久化）
            **kwargs: 额外参数

        Yields:
            流式输出的分析结果，包括：
            - reasoning_step: 推理步骤
            - tool_call: 工具调用
            - final_answer: 最终答案
            - metadata: 元数据
        """
        try:
            # 初始化工具
            self.initialize_tools()

            # 生成线程ID
            if not thread_id:
                thread_id = f"{self.name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

            # 构建输入消息
            messages = self._build_input_messages(input_data, **kwargs)

            # 配置
            config = RunnableConfig(
                configurable={"thread_id": thread_id},
                recursion_limit=50,
            )

            # 流式执行ReAct智能体（使用updates模式获取增量更新）
            reasoning_steps = []
            tool_calls = []
            seen_messages = set()  # 跟踪已处理的消息

            self.logger.info(f"开始流式执行智能体，线程ID: {thread_id}")

            async for event in self.react_agent.astream(
                {"messages": messages},
                config=config,
                stream_mode="updates"  # ✅ 使用updates模式获取增量更新
            ):
                self.logger.debug(f"收到事件: {list(event.keys())}")
                # 处理事件 - updates模式返回节点更新
                for node_name, node_update in event.items():
                    self.logger.debug(f"处理节点更新: {node_name}")
                    if "messages" in node_update:
                        messages_update = node_update["messages"]

                        # 处理新增的消息
                        for msg in messages_update:
                            msg_id = id(msg)
                            if msg_id in seen_messages:
                                continue
                            seen_messages.add(msg_id)

                            # AI消息（推理步骤和工具调用）
                            if isinstance(msg, AIMessage):
                                self.logger.debug(f"处理AI消息，内容长度: {len(msg.content) if msg.content else 0}")
                                # 1. 提取推理内容（如果有）
                                if msg.content and isinstance(msg.content, str) and msg.content.strip():
                                    # 格式化推理内容
                                    reasoning_text = self._format_reasoning_content(msg.content)
                                    if reasoning_text:
                                        self.logger.debug(f"生成推理步骤: {reasoning_text[:100]}...")
                                        reasoning_step = {
                                            "type": "reasoning",
                                            "content": reasoning_text,
                                            "timestamp": datetime.now().isoformat(),
                                        }
                                        reasoning_steps.append(reasoning_text)
                                        yield reasoning_step

                                # 2. 提取工具调用（如果有）
                                if hasattr(msg, 'tool_calls') and msg.tool_calls:
                                    self.logger.info(f"检测到工具调用: {len(msg.tool_calls)}个")
                                    for tool_call in msg.tool_calls:
                                        tool_name = tool_call.get('name', 'unknown')
                                        tool_args = tool_call.get('args', {})

                                        self.logger.info(f"工具调用详情: {tool_name} with args: {tool_args}")

                                        # 生成友好的工具调用描述
                                        tool_desc = self._format_tool_call(tool_name, tool_args)

                                        tool_call_info = {
                                            "type": "reasoning",
                                            "content": tool_desc,
                                            "timestamp": datetime.now().isoformat(),
                                        }
                                        tool_calls.append(tool_call_info)
                                        # ❌ 隐藏工具调用显示 - 不在Step中显示工具调用
                                        # yield tool_call_info
            
            # 获取最终状态
            final_state = await self._get_final_state(thread_id)
            
            # 提取最终答案
            final_answer = self._extract_final_answer(final_state)
            
            # 输出最终结果
            yield {
                "type": "final_answer",
                "content": final_answer,
                "reasoning_steps": reasoning_steps,
                "tool_calls": tool_calls,
                "metadata": {
                    "agent_name": self.name,
                    "thread_id": thread_id,
                    "timestamp": datetime.now().isoformat(),
                },
            }
            
        except Exception as e:
            self.logger.error(f"流式分析失败: {e}")
            yield {
                "type": "error",
                "error": str(e),
                "timestamp": datetime.now().isoformat(),
            }
    
    async def analyze(self, input_data: Any, **kwargs) -> Dict[str, Any]:
        """
        同步分析方法（便捷方法，收集所有流式输出）
        
        Args:
            input_data: 输入数据
            **kwargs: 额外参数
            
        Returns:
            完整的分析结果
        """
        reasoning_steps = []
        tool_calls = []
        final_answer = None
        metadata = {}
        
        async for event in self.stream_analyze(input_data, **kwargs):
            event_type = event.get("type")
            
            if event_type == "reasoning":
                reasoning_steps.append(event.get("content"))
            elif event_type == "tool_call":
                tool_calls.append(event)
            elif event_type == "final_answer":
                final_answer = event.get("content")
                metadata = event.get("metadata", {})
            elif event_type == "error":
                return self.handle_error(Exception(event.get("error")), input_data)
        
        return self.format_response({
            "final_answer": final_answer,
            "reasoning_steps": reasoning_steps,
            "tool_calls": tool_calls,
            "metadata": metadata,
        })
    
    def _build_input_messages(self, input_data: Any, **kwargs) -> List[BaseMessage]:
        """构建输入消息"""
        messages: List[BaseMessage] = [
            SystemMessage(content=self.get_system_prompt())
        ]
        
        # 添加用户消息
        if isinstance(input_data, str):
            messages.append(HumanMessage(content=input_data))
        elif isinstance(input_data, dict):
            # 如果是字典，转换为描述性文本
            content = self._dict_to_text(input_data)
            messages.append(HumanMessage(content=content))
        else:
            messages.append(HumanMessage(content=str(input_data)))
        
        return messages
    
    def _dict_to_text(self, data: Dict[str, Any]) -> str:
        """将字典转换为文本描述"""
        lines = []
        for key, value in data.items():
            if isinstance(value, dict):
                lines.append(f"{key}:")
                for sub_key, sub_value in value.items():
                    lines.append(f"  - {sub_key}: {sub_value}")
            elif isinstance(value, list):
                lines.append(f"{key}: {', '.join(map(str, value))}")
            else:
                lines.append(f"{key}: {value}")
        return "\n".join(lines)
    
    async def _get_final_state(self, thread_id: str) -> Any:
        """获取最终状态"""
        try:
            config = RunnableConfig(configurable={"thread_id": thread_id})
            state = await self.react_agent.aget_state(config)
            return state
        except Exception as e:
            self.logger.warning(f"获取最终状态失败: {e}")
            return None
    
    def _format_reasoning_content(self, content: str) -> str:
        """
        格式化推理内容，使其更适合在Chainlit Step中显示
        
        Args:
            content: 原始推理内容
            
        Returns:
            格式化后的推理内容
        """
        if not content or not isinstance(content, str):
            return ""
        
        content = content.strip()
        
        # ✅ 放宽过滤条件 - 只过滤真正无用的内容
        # 1. 空内容
        if not content:
            return ""
        
        # 2. 太短的内容（放宽限制：从10改为3）
        if len(content) < 3:
            return ""
        
        # ✅ 保留所有推理内容，包括英文思考过程
        # 添加换行符使其更易读
        if not content.endswith('\n'):
            content += '\n'
        
        return content
    
    def _format_tool_call(self, tool_name: str, tool_args: Dict[str, Any]) -> str:
        """
        格式化工具调用，生成友好的描述
        
        Args:
            tool_name: 工具名称
            tool_args: 工具参数
            
        Returns:
            格式化后的工具调用描述
        """
        # 工具名称映射（中文描述）
        tool_descriptions = {
            "get_bond_info": "📋 获取债券基本信息",
            "get_underlying_stock_code": "🔍 查询正股代码",
            "analyze_bond_rating": "⭐ 分析债券评级",
            "get_bond_terms": "📜 获取债券条款",
            "analyze_redemption_terms": "💰 分析赎回条款",
            "get_macro_data": "🌍 获取宏观数据",
            "analyze_industry": "🏭 分析行业情况",
            "get_stock_industry": "🔎 识别正股行业",
            "calculate_valuation": "💎 计算估值指标",
            "compare_similar_bonds": "⚖️ 对比同类债券",
        }
        
        desc = tool_descriptions.get(tool_name, f"🔧 调用工具: {tool_name}")
        
        # 如果有债券代码参数，添加到描述中
        if 'bond_code' in tool_args:
            desc += f" ({tool_args['bond_code']})"
        elif 'stock_code' in tool_args:
            desc += f" ({tool_args['stock_code']})"
        
        return desc + "\n"
    
    def _extract_final_answer(self, state: Any) -> str:
        """从状态中提取最终答案"""
        try:
            if state and hasattr(state, 'values') and "messages" in state.values:
                messages = state.values["messages"]
                if messages:
                    last_message = messages[-1]
                    if isinstance(last_message, AIMessage):
                        content = last_message.content
                        # 处理content可能是list的情况
                        if isinstance(content, list):
                            return ' '.join([str(item) for item in content])
                        return str(content)
            return "分析完成"
        except Exception as e:
            self.logger.warning(f"提取最终答案失败: {e}")
            return "分析完成"
    
    def format_response(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """格式化响应"""
        return {
            "agent": self.name,
            "status": "success",
            "timestamp": datetime.now().isoformat(),
            "data": result
        }
    
    def handle_error(self, error: Exception, input_data: Any, context: str = "") -> Dict[str, Any]:
        """错误处理方法"""
        error_msg = f"{self.name}智能体处理{context}时出错: {str(error)}"
        self.logger.error(error_msg)
        
        return {
            "agent": self.name,
            "status": "error",
            "error": error_msg,
            "input_data": str(input_data),
            "timestamp": datetime.now().isoformat()
        }
    
    def get_capabilities(self) -> Dict[str, Any]:
        """获取智能体能力描述"""
        return {
            "name": self.name,
            "description": self.description,
            "type": self.__class__.__name__,
            "supports_streaming": True,
            "supports_tools": len(self._tools) > 0,
            "num_tools": len(self._tools),
            "version": "2.0.0"
        }
    
    def __str__(self):
        return f"{self.name}智能体({self.description})"
    
    def __repr__(self):
        return f"<{self.__class__.__name__} name='{self.name}' tools={len(self._tools)}>"


class FinancialAnalysisAgent(BaseAgent):
    """金融分析智能体基类 - 统一债券和股票分析"""

    def __init__(self, name: str, description: str = "", llm_config: Optional[LLMConfig] = None):
        super().__init__(name, description, llm_config)
        self.analysis_dimensions = []  # 分析维度列表
        self.asset_type = "financial"  # 资产类型：bond/stock/financial

    def get_system_prompt(self) -> str:
        """获取金融分析的系统提示词"""
        asset_desc = {
            "bond": "债券",
            "stock": "股票",
            "financial": "金融产品"
        }.get(self.asset_type, "金融产品")

        return f"""你是一个专业的{self.name}智能体，负责{asset_desc}{self.description}。

你的职责：
1. 使用提供的工具获取和处理{asset_desc}数据
2. 进行深入的数据分析和推理
3. 提供专业的分析结论和建议

分析维度：{', '.join(self.analysis_dimensions) if self.analysis_dimensions else '综合分析'}

请按照以下步骤进行分析：
1. 理解分析需求
2. 使用工具获取必要数据
3. 分析数据并提取关键信息
4. 得出结论并提供建议

在分析过程中：
- 展示你的推理过程
- 解释每个步骤的目的
- 基于数据得出结论
- 提供可操作的建议
"""

    async def analyze_asset(self, asset_code: str, **kwargs) -> Dict[str, Any]:
        """分析金融资产的通用方法"""
        return await self.analyze(asset_code, **kwargs)

    # 保持向后兼容的方法
    async def analyze_bond(self, bond_code: str, **kwargs) -> Dict[str, Any]:
        """分析债券的便捷方法（向后兼容）"""
        return await self.analyze_asset(bond_code, **kwargs)


# 向后兼容的别名
BondAnalysisAgent = FinancialAnalysisAgent